package ch.android.reasoning;


import java.util.Arrays;
import java.util.List;

/**
 * Original
 * @author Ravi Mohan
 * @version 1.0
 * @date 17.02.2005
 * 
 * Modification
 * @author voigt
 * @version 1.0.1
 * @date 22.05.2008
 * 
 */

public class HiddenMarkovModel {

	private RandomVariable priorDistribution;
	private SensorModel sensorModel;
	private TransitionModel transitionModel;
	
	// Getters
	public SensorModel getSensorModel() {

		return sensorModel;
	}

	public TransitionModel getTransitionModel() {

		return transitionModel;
	}
	
	public RandomVariable getPriorDistribution() {
		return priorDistribution;
	}

	public HiddenMarkovModel()
	{
	}
	
	public HiddenMarkovModel(RandomVariable priorDistribution, TransitionModel tm, SensorModel sm) {
		this.priorDistribution = priorDistribution;
		this.transitionModel = tm;
		this.sensorModel = sm;
	}

	// Calculates T^T*f_1:t
	private RandomVariable predict(RandomVariable aBelief, String action) {
		RandomVariable newBelief = aBelief.duplicate();
		Matrix beliefMatrix = aBelief.asMatrix();
		Matrix transitionMatrix = transitionModel.asMatrix(action);
		Matrix predicted = transitionMatrix.transpose().times(beliefMatrix);
		newBelief.updateFrom(predicted);
		
		return newBelief;
	}

	// Calculates O_t+1 and multiplies it with (T^t*f_1:t) from predict() 
	// in case of the forward pass.
	private RandomVariable perceptionUpdate(RandomVariable aBelief, String perception) {
		RandomVariable newBelief = aBelief.duplicate();
		// one way - use matrices
		Matrix beliefMatrix = aBelief.asMatrix();
		Matrix observationMatrix = sensorModel.asMatrix(perception);
		Matrix updated = observationMatrix.times(beliefMatrix);
		newBelief.updateFrom(updated);
		newBelief.normalize();
		
		return newBelief;
		
	}

	private RandomVariable forward(RandomVariable aBelief, String action,
			String perception) {

		return perceptionUpdate(predict(aBelief, action), perception);
	}

	private RandomVariable forward(RandomVariable aBelief, String perception) {

		return forward(aBelief, HmmConstants.DO_NOTHING, perception);
	}

	// Calculates b = T*O_k+1+b_k+2:t _each time_ its called 
	// where:
	// b = result
	// T = transitionMatrix
	// O_k+1 = observationMatrix
	// b_k+2:t = presentBackwardBeliefe
	private RandomVariable backward(RandomVariable forwardBelief, RandomVariable presentBackwardBelief, String perception) {
		
		RandomVariable result = presentBackwardBelief.duplicate();
		Matrix observationMatrix = sensorModel.asMatrix(perception);
		Matrix transitionMatrix = transitionModel.asMatrix();
		Matrix backwardMatrix = transitionMatrix.times(observationMatrix.times(presentBackwardBelief.asMatrix()));
		Matrix resultMatrix = backwardMatrix.arrayTimes(forwardBelief.asMatrix());
		result.updateFrom(resultMatrix);
		result.normalize();
		
		return result;
	}

	// TODO: save on the .asMatrix() calls by either using
	// Matrix as native type or store ever forward/backward message
	// in a redundant second Matrix type variable.
	public List<RandomVariable> forward_backward(List<String> perceptions) {
		RandomVariable forwardMessages[] = new RandomVariable[perceptions.size() + 1];
		RandomVariable backwardMessage = priorDistribution.createUnitBelief();
		RandomVariable smoothedBeliefs[] = new RandomVariable[perceptions.size() + 1];

		forwardMessages[0] = priorDistribution;
		smoothedBeliefs[0] = null;

		// populate forward messages for all perceptions
		for (int i = 0; i < perceptions.size(); i++) { // N.B i starts at 1, not zero
			forwardMessages[i + 1] = forward(forwardMessages[i], perceptions.get(i));
		}
		for (int i = perceptions.size(); i > 0; i--) {
			RandomVariable smoothed = priorDistribution.duplicate();

			// Calculate element multiplication of forward and backward messages, then normalize
			smoothed.updateFrom(forwardMessages[i].asMatrix().arrayTimes(backwardMessage.asMatrix()));
			smoothed.normalize();
			smoothedBeliefs[i] = smoothed;
			
			// Go back one more perception and calculate its backward message
			backwardMessage = backward(forwardMessages[i], backwardMessage, perceptions.get(i - 1));
		}

		return Arrays.asList(smoothedBeliefs);
	}

}
